from flask import Flask, render_template, request, jsonify, session, redirect, url_for
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
import io
import base64
import json
import os
from datetime import datetime, timedelta

app = Flask(__name__)
app.secret_key = 'rehabilitation_web_app_secret_key'

# 模拟数据库
class RehabilitationDatabase:
    def __init__(self):
        self.patients = {}
        self.doctors = {}
        self.rehabilitation_records = {}
        self.predictions = {}
        self.next_patient_id = 1
        self.next_doctor_id = 1
        self.next_record_id = 1
    
    def add_patient(self, patient_data):
        patient_id = self.next_patient_id
        patient_data['id'] = patient_id
        patient_data['registration_date'] = datetime.now().strftime('%Y-%m-%d')
        self.patients[patient_id] = patient_data
        self.next_patient_id += 1
        return patient_id
    
    def add_doctor(self, doctor_data):
        doctor_id = self.next_doctor_id
        doctor_data['id'] = doctor_id
        self.doctors[doctor_id] = doctor_data
        self.next_doctor_id += 1
        return doctor_id
    
    def add_rehabilitation_record(self, record_data):
        record_id = self.next_record_id
        record_data['id'] = record_id
        record_data['record_date'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.rehabilitation_records[record_id] = record_data
        self.next_record_id += 1
        return record_id
    
    def get_patient_records(self, patient_id):
        return [record for record in self.rehabilitation_records.values() 
                if record.get('patient_id') == patient_id]
    
    def get_doctor_patients(self, doctor_id):
        return [patient for patient in self.patients.values() 
                if patient.get('assigned_doctor') == doctor_id]

# 初始化数据库
db = RehabilitationDatabase()

# 康复预测模型
class RehabilitationPredictor:
    def __init__(self):
        self.model = RandomForestRegressor(n_estimators=100, random_state=42)
        self.scaler = StandardScaler()
        self.is_trained = False
    
    def generate_sample_data(self, n_samples=1000):
        """生成模拟康复数据"""
        np.random.seed(42)
        
        data = {
            'age': np.random.randint(18, 80, n_samples),
            'gender': np.random.choice([0, 1], n_samples),  # 0: 女, 1: 男
            'injury_type': np.random.choice([0, 1, 2, 3], n_samples),  # 损伤类型
            'injury_severity': np.random.randint(1, 10, n_samples),  # 损伤严重程度
            'treatment_type': np.random.choice([0, 1, 2], n_samples),  # 治疗类型
            'compliance_rate': np.random.uniform(0.5, 1.0, n_samples),  # 依从率
            'initial_mobility': np.random.uniform(0, 100, n_samples),  # 初始活动能力
            'pain_level': np.random.randint(1, 10, n_samples),  # 疼痛程度
            'rehabilitation_days': np.random.randint(30, 365, n_samples)  # 康复天数
        }
        
        return pd.DataFrame(data)
    
    def train_model(self):
        """训练预测模型"""
        df = self.generate_sample_data(1000)
        
        # 特征和目标变量
        features = ['age', 'gender', 'injury_type', 'injury_severity', 
                   'treatment_type', 'compliance_rate', 'initial_mobility', 'pain_level']
        X = df[features]
        y = df['rehabilitation_days']
        
        # 数据标准化
        X_scaled = self.scaler.fit_transform(X)
        
        # 训练模型
        self.model.fit(X_scaled, y)
        self.is_trained = True
        
        return "模型训练完成"
    
    def predict_rehabilitation(self, patient_data):
        """预测康复周期"""
        if not self.is_trained:
            self.train_model()
        
        # 准备特征数据
        features = ['age', 'gender', 'injury_type', 'injury_severity', 
                   'treatment_type', 'compliance_rate', 'initial_mobility', 'pain_level']
        
        # 确保所有特征都存在
        for feature in features:
            if feature not in patient_data:
                patient_data[feature] = 0
        
        X = pd.DataFrame([patient_data])[features]
        X_scaled = self.scaler.transform(X)
        
        # 预测
        predicted_days = int(self.model.predict(X_scaled)[0])
        
        # 风险评估
        risk_level = self.assess_risk(patient_data, predicted_days)
        
        return {
            'predicted_days': predicted_days,
            'risk_level': risk_level,
            'estimated_completion_date': self.calculate_completion_date(predicted_days)
        }
    
    def assess_risk(self, patient_data, predicted_days):
        """评估康复风险"""
        risk_score = 0
        
        # 年龄风险
        if patient_data['age'] > 65:
            risk_score += 2
        elif patient_data['age'] > 50:
            risk_score += 1
        
        # 损伤严重程度风险
        risk_score += patient_data['injury_severity'] * 0.5
        
        # 疼痛程度风险
        risk_score += patient_data['pain_level'] * 0.3
        
        # 依从率风险
        if patient_data['compliance_rate'] < 0.7:
            risk_score += 2
        
        if risk_score < 3:
            return "低风险"
        elif risk_score < 6:
            return "中风险"
        else:
            return "高风险"
    
    def calculate_completion_date(self, days):
        """计算预计完成日期"""
        return (datetime.now() + timedelta(days=days)).strftime('%Y-%m-%d')

# 初始化预测器
predictor = RehabilitationPredictor()

# 路由定义
@app.route('/')
def index():
    """首页 - 选择登录类型"""
    return render_template('index.html')

@app.route('/doctor')
def doctor_dashboard():
    """医护端主页"""
    return render_template('doctor_dashboard.html')

@app.route('/patient')
def patient_dashboard():
    """患者端主页"""
    return render_template('patient_dashboard.html')

@app.route('/api/predict', methods=['POST'])
def api_predict():
    """康复周期预测API"""
    try:
        data = request.json
        prediction = predictor.predict_rehabilitation(data)
        return jsonify({'success': True, 'prediction': prediction})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/api/patients', methods=['GET'])
def api_get_patients():
    """获取患者列表"""
    patients = list(db.patients.values())
    return jsonify({'success': True, 'patients': patients})

@app.route('/api/patient/<int:patient_id>', methods=['GET'])
def api_get_patient(patient_id):
    """获取特定患者信息"""
    patient = db.patients.get(patient_id)
    if patient:
        records = db.get_patient_records(patient_id)
        return jsonify({'success': True, 'patient': patient, 'records': records})
    else:
        return jsonify({'success': False, 'error': '患者不存在'})

@app.route('/api/patient/add', methods=['POST'])
def api_add_patient():
    """添加新患者"""
    try:
        patient_data = request.json
        patient_id = db.add_patient(patient_data)
        return jsonify({'success': True, 'patient_id': patient_id})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/api/record/add', methods=['POST'])
def api_add_record():
    """添加康复记录"""
    try:
        record_data = request.json
        record_id = db.add_rehabilitation_record(record_data)
        return jsonify({'success': True, 'record_id': record_id})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/api/visualization/patient_progress/<int:patient_id>')
def api_patient_progress_visualization(patient_id):
    """生成患者康复进度可视化"""
    records = db.get_patient_records(patient_id)
    if not records:
        return jsonify({'success': False, 'error': '无康复记录'})
    
    # 创建进度图表
    dates = [record['record_date'] for record in records]
    mobility_scores = [record.get('mobility_score', 0) for record in records]
    pain_levels = [record.get('pain_level', 0) for record in records]
    
    plt.figure(figsize=(10, 6))
    plt.plot(dates, mobility_scores, marker='o', label='活动能力')
    plt.plot(dates, pain_levels, marker='s', label='疼痛程度')
    plt.xlabel('日期')
    plt.ylabel('评分')
    plt.title('康复进度跟踪')
    plt.legend()
    plt.xticks(rotation=45)
    plt.tight_layout()
    
    # 转换为base64图像
    img = io.BytesIO()
    plt.savefig(img, format='png')
    img.seek(0)
    plot_url = base64.b64encode(img.getvalue()).decode()
    plt.close()
    
    return jsonify({'success': True, 'plot_url': f"data:image/png;base64,{plot_url}"})

if __name__ == '__main__':
    # 预训练模型
    predictor.train_model()
    
    # 添加示例数据
    sample_patient = {
        'name': '张三',
        'age': 45,
        'gender': 1,
        'injury_type': '骨折',
        'injury_severity': 7,
        'contact': '13800138000',
        'assigned_doctor': 1
    }
    db.add_patient(sample_patient)
    
    sample_doctor = {
        'name': '李医生',
        'specialty': '康复医学',
        'contact': '13900139000'
    }
    db.add_doctor(sample_doctor)
    
    app.run(debug=True, host='0.0.0.0', port=5000)